from typing import Optional, Dict, Any
import os
import uvicorn
from fastapi import FastAPI, HTTPException, Body, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
import threading
import time
from loguru import logger
from utils.logger import Logger
from utils.constants import AgentType, AGENT_MAX_CONCURRENT
from states.agent_state import AgentState
from middleware.redis_adapter import get_redis_adapter

# 文档处理依赖（需安装：pip install langchain pdfplumber python-docx）
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.embeddings.openai import OpenAIEmbeddings
import pdfplumber
from docx import Document


# ===================== 文档加载与向量库管理 =====================
class DocVectorStore:
    def __init__(self, openai_api_key: str):
        self.embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
        self.vector_store = None
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50,
            length_function=len
        )

    def load_document(self, file_path: str) -> str:
        """加载PDF/Word文档文本"""
        ext = os.path.splitext(file_path)[1].lower()
        text = ""
        try:
            if ext == ".pdf":
                with pdfplumber.open(file_path) as pdf:
                    for page in pdf.pages:
                        text += page.extract_text() or ""
            elif ext == ".docx":
                doc = Document(file_path)
                for paragraph in doc.paragraphs:
                    text += paragraph.text + "\n"
            else:
                raise ValueError(f"不支持的文档格式：{ext}，仅支持PDF/Word")
            return text.strip()
        except Exception as e:
            logger.error(f"文档加载失败：{str(e)}")
            raise RuntimeError(f"文档处理异常：{str(e)}")

    def build_vector_store(self, text: str) -> None:
        """构建向量库（支持后续相似性查询）"""
        chunks = self.text_splitter.split_text(text)
        self.vector_store = FAISS.from_texts(chunks, self.embeddings)
        logger.info(f"向量库构建完成，文本分片数：{len(chunks)}")

    def similarity_search(self, query: str, k: int = 3) -> str:
        """相似性查询，获取相关上下文"""
        if not self.vector_store:
            return "未加载任何文档，请先上传文档并构建向量库"
        docs = self.vector_store.similarity_search(query, k=k)
        return "\n\n".join([doc.page_content for doc in docs])


# ===================== 文档问答Agent核心服务 =====================
class DocQAAgentService:
    def __init__(self):
        # 环境变量配置
        self.host = os.getenv("DOC_QA_HOST", "0.0.0.0")
        self.port = int(os.getenv("DOC_QA_PORT", 8004))  # 独立端口
        self.openai_api_key = os.getenv("OPENAI_API_KEY", "")
        self.doc_storage_path = os.path.join(os.getcwd(), "apps", "doc_qa", "uploaded_docs")

        # 创建文档存储目录
        if not os.path.exists(self.doc_storage_path):
            os.makedirs(self.doc_storage_path)

        # 核心组件初始化
        self.redis_client = get_redis_adapter()  # 使用Redis适配器，支持集群模式
        self.logger = Logger.get_logger(agent_name="doc_qa_agent", task_id="none")
        self.doc_vector_store = DocVectorStore(self.openai_api_key)  # 文档向量库
        self.current_doc_name = None  # 当前加载的文档名称

        # AgentState注册
        self.instance_id = f"doc_qa_{self.host}:{self.port}"
        self.agent_state = AgentState(
            agent_type=AgentType.DOC_QA.value,
            instance_id=self.instance_id,
            redis_client=self.redis_client,
            host_ip=self.host,
            port=self.port
        )

        # 启动心跳
        self._start_heartbeat()

        # FastAPI服务初始化
        self.app = self._init_fastapi()
        self.logger.success(
            f"文档问答Agent启动完成 | 实例ID：{self.instance_id} | 地址：http://{self.host}:{self.port} | 文档存储路径：{self.doc_storage_path}")

    def _init_fastapi(self) -> FastAPI:
        app = FastAPI(title="Document QA Agent", version="v1.0")

        # 跨域中间件
        app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

        # 健康检查接口
        @app.get("/health")
        async def health_check():
            self.logger.debug(f"健康检查请求 | 实例：{self.instance_id}")
            return {
                "status": self.agent_state.status,
                "instance_id": self.instance_id,
                "cpu_usage": self.agent_state.cpu_usage,
                "mem_usage": self.agent_state.mem_usage,
                "current_tasks": self.agent_state.current_tasks,
                "loaded_doc": self.current_doc_name or "无"
            }

        # 文档上传接口
        @app.post("/upload_doc", summary="上传文档（PDF/Word）")
        async def upload_doc(
                file: UploadFile = File(..., description="PDF/Word文档"),
                task_id: Optional[str] = Body("none", description="任务ID")
        ):
            task_logger = Logger.update_task_id(self.logger, task_id)
            task_logger.info(f"收到文档上传请求 | 文件名：{file.filename}")
            try:
                # 保存文档到本地
                file_ext = os.path.splitext(file.filename)[1].lower()
                if file_ext not in [".pdf", ".docx"]:
                    raise ValueError("仅支持PDF和Word文档（.pdf/.docx）")

                file_path = os.path.join(self.doc_storage_path, file.filename)
                with open(file_path, "wb") as f:
                    f.write(await file.read())

                # 加载文档并构建向量库
                doc_text = self.doc_vector_store.load_document(file_path)
                self.doc_vector_store.build_vector_store(doc_text)
                self.current_doc_name = file.filename

                self.agent_state.update_task_metrics(success=True)
                task_logger.success(f"文档上传并构建向量库成功 | 文件名：{file.filename} | 文本长度：{len(doc_text)}字")
                return {
                    "success": True,
                    "data": {
                        "filename": file.filename,
                        "file_path": file_path,
                        "text_length": len(doc_text),
                        "message": "文档已加载，可进行问答"
                    },
                    "error_code": 0,
                    "error_msg": ""
                }
            except Exception as e:
                error_msg = str(e)
                self.agent_state.update_task_metrics(success=False)
                task_logger.error(f"文档上传失败：{error_msg}")
                raise HTTPException(
                    status_code=500,
                    detail={"success": False, "error_code": 2002, "error_msg": error_msg}
                )

        # 文档问答核心接口（和text_qa统一格式）
        @app.post("/predict")
        async def predict(
                question: str = Body(..., description="针对文档的问题"),
                task_id: Optional[str] = Body("none", description="任务ID")
        ):
            task_logger = Logger.update_task_id(self.logger, task_id)
            task_logger.info(f"收到文档问答请求 | 问题：{question[:30]}...")
            try:
                if not self.current_doc_name:
                    raise ValueError("未加载任何文档，请先通过/upload_doc接口上传PDF/Word文档")

                # 从向量库获取相关上下文
                relevant_context = self.doc_vector_store.similarity_search(question)
                task_logger.debug(f"匹配到的文档上下文：{relevant_context[:100]}...")

                # 调用OpenAI生成答案（基于文档上下文）
                from openai import OpenAI
                client = OpenAI(api_key=self.openai_api_key)
                response = client.chat.completions.create(
                    model="gpt-3.5-turbo",
                    messages=[
                        {"role": "system",
                         "content": "严格基于提供的文档上下文回答问题，不编造信息；如果上下文没有相关内容，直接说'文档中未找到相关答案'"},
                        {"role": "user", "content": f"文档上下文：{relevant_context}\n问题：{question}"}
                    ],
                    temperature=0.2,
                    max_tokens=500
                )
                answer = response.choices[0].message.content.strip()

                self.agent_state.update_task_metrics(success=True)
                task_logger.info(f"文档问答处理成功 | 答案：{answer[:30]}...")
                return {
                    "success": True,
                    "data": {
                        "question": question,
                        "answer": answer,
                        "loaded_doc": self.current_doc_name,
                        "instance_id": self.instance_id
                    },
                    "error_code": 0,
                    "error_msg": ""
                }
            except Exception as e:
                error_msg = str(e)
                self.agent_state.update_task_metrics(success=False)
                task_logger.error(f"文档问答处理失败：{error_msg}")
                raise HTTPException(
                    status_code=500,
                    detail={"success": False, "error_code": 2002, "error_msg": error_msg}
                )

        return app

    def _start_heartbeat(self):
        """心跳线程（和其他Agent统一逻辑）"""

        def heartbeat():
            while True:
                try:
                    self.agent_state.send_heartbeat()
                    self.logger.debug(
                        f"心跳同步成功 | 状态：{self.agent_state.status} | 任务数：{self.agent_state.current_tasks} | 已加载文档：{self.current_doc_name or '无'}")
                except Exception as e:
                    self.logger.error(f"心跳同步失败：{str(e)}")
                time.sleep(10)

        thread = threading.Thread(target=heartbeat, daemon=True)
        thread.start()
        self.logger.info("心跳线程启动成功")

    def run(self):
        uvicorn.run(app=self.app, host=self.host, port=self.port, log_level="warning")


# 启动入口
def main():
    service = DocQAAgentService()
    service.run()


if __name__ == "__main__":
    main()